import os
import json

from config import label_save_path, predict_dir
from utils.bbox import IOU

class Benchmark:
    def __init__(self, labelpath, predict_dir):
        self.labelpath = labelpath
        self.predict_dir = predict_dir
        self.label = None
        self.predict = dict()
        self.match_result_refer_predict = dict()
        self.match_result_refer_gt = dict()
        self.loadLabel()
        self.loadPredict()
                
    def loadLabel(self):
        with open(self.labelpath, "r") as fp:
            label = json.load(fp)
        
        new_label = dict()
        for key in label.keys():
            value = label[key]
            refer = value[0]
            subfigures = value[1]
            localPositionSubfigures = []
            for subfigure in subfigures:
                localPositionSubfigures.append(self.convertLocalPosition(refer, subfigure))
            new_label[key] = localPositionSubfigures
        self.label = new_label
    
    @staticmethod
    def convertLocalPosition(refer, obj):
        return [
            obj[0] - refer[0],
            obj[1] - refer[1],
            obj[2] - refer[0],
            obj[3] - refer[1],
        ]
        
                       
    def loadPredict(self):
        for filename in os.listdir(self.predict_dir):
            if not filename.endswith("json"):
                continue
            filepath = os.path.join(self.predict_dir, filename)
            with open(filepath , "r") as fp:
                data = json.load(fp)
            
            bboxes = []
            for item in data:
                l = item["x"]
                t = item["y"]
                r = item["w"] + l
                b = item["h"] + t
                bbox = [l, t, r, b]
                bboxes.append(bbox)
            
            # key 
            imagename = filename[0:len(filename)-5]
            self.predict[imagename] = bboxes
            
        
    def run(self, threshold):
        for key in self.predict.keys():
            predicts = self.predict[key]
            labels = self.label[key]
            predict_match, label_match = self.dealOne(labels, predicts)
            self.match_result_refer_gt[key] = label_match
            self.match_result_refer_predict[key] = predict_match
        
        self.getScore(threshold)    
        
    def getScore(self, threshold):
        # 计算mIOU
        IOU = []
        SFCR = []
        totalIOU = 0
        n = 0
        for key in self.match_result_refer_gt.keys():
            item = self.match_result_refer_gt[key]
            IOU_item = []
            recgnized = 0
            for match in item:
                iou = match[-1]
                IOU_item.append(iou)
                totalIOU += iou
                n += 1
                if iou > threshold:
                    recgnized += 1
            SFCR.append([len(item), recgnized])
            IOU.append(IOU_item)
        
        mIOU = totalIOU / n
        
        totalSFCR = 0         
        for item in SFCR:
            totalSFCR += (item[1] / item[0])
        
        mSFCR = totalSFCR / len(SFCR)
        totalSubfigures, recognizedSubfigures = sum([item[0] for item in SFCR]), sum([item[1] for item in SFCR])
        print(
            f"threshold: {round(threshold, 3)}",
            f"mIOU: {round(mIOU, 3)}",
            f"mSFCR: {round(mSFCR, 3)}",
            f"totalSubfigures: {round(totalSubfigures, 3)}",
            f"recognizedSubfigures: {round(recognizedSubfigures, 3)}"
        )
        return mIOU, mSFCR, totalSubfigures, recognizedSubfigures
                
    def dealOne(self, labels, predicts):
        predict_match = []
        for predict in predicts:
            max_iou = IOU(predict, labels[0])
            max_iou_label = labels[0]
            for idx in range(1, len(labels)):
                label = labels[idx]
                iou = IOU(predict, label)
                if iou > max_iou:
                    max_iou = iou
                    max_iou_label = label
            predict_match.append([max_iou_label, predict, max_iou])
            
        gt_match = []
        for label in labels:
            max_iou = IOU(label, predicts[0])
            max_iou_label = labels[0]
            for idx in range(1, len(predicts)):
                predict = predicts[idx]
                iou = IOU(predict, label)
                if iou > max_iou:
                    max_iou = iou
                    max_iou_label = predict
            gt_match.append([label, max_iou_label, max_iou])
        return predict_match, gt_match
            
if __name__ == "__main__":
    benchmark = Benchmark(label_save_path, predict_dir)
    benchmark.run(0.7)
    for i in range(10, 20, 1):
        benchmark.getScore(round(i*0.05, 2))